import keras.callbacks
import Util as util


class ModelSave(keras.callbacks.Callback):
    def __init__(self, lossSavePath, lossProgressSavePath, modelSavePath, weightsSavePath):
        self.lossSavePath = lossSavePath
        self.lossProgressSavePath = lossProgressSavePath
        self.modelPath = modelSavePath
        self.weightsPath = weightsSavePath

    def on_epoch_end(self, epoch, logs=None):
        lossNew = logs.get('loss')
        try:
            file = open(self.lossSavePath, 'r')
            lossOldStr = file.read()
            file.close()
        except:
            lossOldStr = ''

        if not lossOldStr or lossNew < float(lossOldStr):
            # overwrite
            with open(self.lossSavePath, 'w') as file:
                file.write(str(lossNew))

            # record train loss progress
            if epoch == 0:
                with open(self.lossProgressSavePath, 'w') as file:
                    file.write("#epoch" + str(epoch) + ": " + str(lossNew) + '\n')
            else:
                with open(self.lossProgressSavePath, 'a') as file:
                    file.write("#epoch" + str(epoch) + ": " + str(lossNew) + '\n')

            # save model and weights
            util.saveModel(self.model, self.modelPath)
            self.model.save_weights(self.weightsPath)
